import base64
import cv2
import numpy as np
from multiprocessing import Process, Queue
import time


def torch_worker(model_path, input_queue, output_queue):
    """子进程：运行torch模型
    处理输入可以是文件路径或内存中的cv2图像
    """
    import torch  # 在子进程内导入
    from ultralytics import YOLO  # 在子进程内导入

    model = YOLO(model_path)

    while True:
        task = input_queue.get()
        if task is None:  # 终止信号
            break

        # 根据输入类型获取图像
        if isinstance(task, str):  # 文件路径
            image = cv2.imread(task)
        else:  # 假定是numpy数组
            image = task

        results = model.predict(image, imgsz=640, conf=0.5)
        boxes = [box.xyxy.cpu().numpy() for box in results[0].boxes]

        # # 绘制所有检测框
        # for box_group in boxes:  # 遍历每个检测框组
        #     for box in box_group:  # 遍历每个检测框
        #         x1, y1, x2, y2 = map(int, box)
        #         cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        #
        # # 保存绘制后的图像
        # cv2.imwrite(f"detected_{time.time()}.jpg", image)

        output_queue.put(boxes)


class PlateDetectionService:
    def __init__(self, yolo_model_path):
        self.input_queue = Queue()
        self.output_queue = Queue()

        # 启动子进程（不传递任何Paddle/不可pickle对象）
        self.process = Process(
            target=torch_worker,
            args=(yolo_model_path, self.input_queue, self.output_queue)
        )
        self.process.start()

        # 主进程初始化PaddleOCR
        from PlateRecognizer import PlateRecognizer
        self.recognizer = PlateRecognizer()

    def process_image(self, image_input):
        """主流程：发送任务到子进程并获取结果
        支持输入文件路径或cv2读取的图像（numpy数组）
        """
        # 根据输入类型获取图像
        if isinstance(image_input, str):  # 文件路径
            image = cv2.imread(image_input)
            if image is None:
                raise ValueError("无法读取图像文件")
        elif isinstance(image_input, np.ndarray):  # numpy数组
            image = image_input
        else:
            raise TypeError("输入类型必须是文件路径或numpy数组")

        # 发送任务到子进程
        self.input_queue.put(image)
        boxes = self.output_queue.get()

        results = []
        for boxgroup in boxes:  # 处理第一个检测结果的boxes
            for box in boxgroup:  # 遍历每个检测框
                x1, y1, x2, y2 = map(int, box)
                plate_img = image[y1:y2, x1:x2]

                # 将 plate_img 转换为 Base64 编码的字符串
                _, buffer = cv2.imencode('.jpg', plate_img)
                plate_img_base64 = base64.b64encode(buffer).decode('utf-8')

                text, confidence = self.recognizer.recognize_plate(plate_img)
                results.append((text, confidence, plate_img_base64))

        return results

    def __del__(self):
        """清理子进程"""
        self.input_queue.put(None)
        self.process.join()


if __name__ == "__main__":
    # 示例用法
    service = PlateDetectionService("Yolov8_PlateDetect.pt")
    # 使用内存中的图像
    image = cv2.imread("1743233055747.jpg")
    results = service.process_image(image)

    for text, conf, _ in results:
        print(f"识别结果: {text} (置信度: {conf:.2f})")
